from abc import ABC, abstractmethod
import numpy as np
from tqdm import tqdm
import matplotlib.pyplot as plt
import platgo as pg


class Algorithm(ABC):
    """
    算法模板基类
    """
    def __init__(self, problem: pg.Problem, maxgen: int) -> None:
        self.name = 'Algorithm'
        self.problem = problem
        self.maxgen = maxgen  # 最大进化代数
        self.gen = 0  # 当前进化代数
        self.track = pg.Track()  # 进化追踪器
        self._p_bar = tqdm(total=maxgen)  # 进度条
        self.is_bar = True  # 是否显示进度条
    
    @abstractmethod
    def go(self):
        """算法运行主函数"""
        pass
    
    def _show_bar(self) -> None:
        """更新进度条"""
        if not self.is_bar:
            return
        if self.gen >= self.maxgen:
            self._p_bar.close()
        else:
            self._p_bar.update(1)
    
    def not_terminal(self, pop: pg.Population) -> bool:
        """检查算法是否到结束条件"""
        self._show_bar()
        if self.gen >= self.maxgen:
            return False
        else:
            self.track.objv_his.append(pop.objv)
            self.gen += 1
            return True
    
    def draw(self):
        """
        画图，
        NOTEb 还没想好需要哪些功能
        :return:
        """
        if self.problem.M == 1:
            objv_his = np.asarray(self.track.objv_his)
            median_objv = np.median(objv_his, axis=1)
            plt.plot(median_objv, label='median objective')
            best_objv = np.min(objv_his, axis=1)
            plt.plot(best_objv, label='best objective')
            plt.xlabel('number of generation')
            plt.ylabel('value of objective')
            plt.legend()
            plt.show()
        else:
            raise NotImplementedError("multi-objective drawing has not been completed!")
